﻿#include <ConstructionManager.h>
#include <UnitGroupManager.h>
using namespace BWAPI;
using namespace std;
ConstructionManager::ConstructionManager(Arbitrator::Arbitrator<BWAPI::Unit*,double>* arbitrator, BuildingPlacer* placer)
{
	this->arbitrator   = arbitrator;
	this->placer       = placer;
	for(std::set<BWAPI::UnitType>::const_iterator i=BWAPI::UnitTypes::allUnitTypes().begin();i!=BWAPI::UnitTypes::allUnitTypes().end();i++)
	{
		plannedCount[*i]=0;
		startedCount[*i]=0;
	}
	this->worker = WorkerManager::create();
}

void ConstructionManager::onOffer(std::set<BWAPI::Unit*> units)
{
	//we are being offered some units (hopefully builders).
	//accept want we need and decline the rest

	//go through all the buildings that are asking for builders
	for(std::map<BWAPI::UnitType,std::set<Building*> >::iterator i=this->buildingsNeedingBuilders.begin();i!=this->buildingsNeedingBuilders.end();i++)
	{
		std::set<Building*>::iterator b_next;
		//go through all the buildings that are asking for builders
		for(std::set<Building*>::iterator b = i->second.begin(); b != i->second.end(); b=b_next)
		{
			b_next=b;
			b_next++;
			double min_dist = 1000000;
			BWAPI::Unit* builder = NULL;
			//find the best builder based on distance
			for(std::set<BWAPI::Unit*>::iterator u = units.begin(); u != units.end(); u++)
			{
				//only consider this builder if it can build this type of building
				if (((*b)->type.whatBuilds().first)==(*u)->getType() && (!(*b)->type.isAddon() || (*u)->getAddon()==NULL))
				{
					double dist = (*u)->getPosition().getDistance((*b)->position) +
                        ((*u)->getType().maxHitPoints() - (*u)->getHitPoints())*32;
					if (dist < min_dist)
					{
						min_dist = dist;
						builder  = *u;
					}
				}
			}

			//if we found a builder for this building
			if (builder != NULL)
			{
				//tell the arbitrator we accept the unit, and raise the bid to hopefully prevent other managers from using it
				arbitrator->accept(this, builder);
				arbitrator->setBid(this, builder, 100.0);

				//add builder to builders map
				builders.insert(std::make_pair(builder, *b));

				//also tell the building it has this builder
				(*b)->builderUnit = builder;

				//if this is the addon, we can also set the tile position and reserve the build site
				if ((*b)->type.isAddon())
				{
					(*b)->tilePosition = builder->getTilePosition();
					if (builder->isLifted())
					{
						if (!this->placer->canBuildHereWithSpace((*b)->tilePosition,(*b)->type))
							(*b)->tilePosition=placer->getBuildLocationNear((*b)->tilePosition,(*b)->type.whatBuilds().first);
					}
					else
					{
						bool buildable=true;
						for(int x=(*b)->tilePosition.x()+4;x<(*b)->tilePosition.x()+6;x++)
							for(int y=(*b)->tilePosition.y()+1;y<(*b)->tilePosition.y()+3;y++)
								if (!this->placer->buildable(x,y))
									buildable=false;
						if (!buildable)
						{
							this->placer->freeTiles((*b)->tilePosition, 4,3);
							(*b)->tilePosition=placer->getBuildLocationNear((*b)->tilePosition,(*b)->type.whatBuilds().first);
						}
					}
					this->placer->reserveTiles((*b)->tilePosition, 4,3);
					this->placer->reserveTiles((*b)->tilePosition+BWAPI::TilePosition(4,1), 2,2);
				}

				//remove builder from the units set
				units.erase(builder);

				//erase this building from the set of buildings that need builders
				i->second.erase(b);
			}
		}
	}
	//decline whatever is left
	for(std::set<BWAPI::Unit*>::iterator u = units.begin(); u != units.end(); u++)
	{
		arbitrator->decline(this, *u, 0);
		arbitrator->removeBid(this, *u);
	}
}

void ConstructionManager::onRevoke(BWAPI::Unit* unit, double bid)
{
	this->onRemoveUnit(unit);
}

void ConstructionManager::update()
{
	std::set<BWAPI::Unit*> myPlayerUnits = SelectAll()(isCompleted)(GetAddon,(BWAPI::Unit*)NULL).not(isCarryingMinerals,isCarryingGas,isGatheringGas);

	//iterate through all the builder types
	for(std::map<BWAPI::UnitType,std::set<Building*> >::iterator i=this->buildingsNeedingBuilders.begin();i!=this->buildingsNeedingBuilders.end();i++)
	{
		if (!i->second.empty()) //if its empty, we don't really need builders of this type
		{
			for(std::set<BWAPI::Unit*>::iterator u = myPlayerUnits.begin(); u != myPlayerUnits.end(); u++)
			{
				// if this unit is completed and the right type, and doesn't have an addon, and we aren't already using it
        // also we make sure that builder has enough HP
				if ((*u)->getType()==i->first && this->builders.find(*u)==this->builders.end())
				{
					//bid value depends on distance - we like closer builders better
					double min_dist=1000000;
					for(std::set<Building*>::iterator b = i->second.begin(); b != i->second.end(); b++)
					{
            // we consider HP into choosing unit too
						double dist = (*u)->getPosition().getDistance((*b)->position) +
                          ((*u)->getType().maxHitPoints() - (*u)->getHitPoints())*32;
						if (dist < min_dist)
							min_dist = dist;
					}
					if (min_dist < 10)
						min_dist = 10;
					if (min_dist > 256*32 + 10)
						min_dist = 256*32 + 10;
					double bid = 80 - (min_dist - 10)/(256*32)*60;
					//if this isn't a worker, we don't care about distance
					if (!(*u)->getType().isWorker())
						arbitrator->setBid(this, *u, 80);
					else
						arbitrator->setBid(this, *u, bid);
				}
			}
		}
	}
	std::list<Building>::iterator i_next;
	int index=0;

	//iterate through all the incomplete buildings
	for(std::list<Building>::iterator i=this->incompleteBuildings.begin();i!=this->incompleteBuildings.end();i=i_next)
	{
		index++;
		i_next=i;
		i_next++;
		Building* b = &(*i);
		if (!b->started && b->buildingUnit!=NULL) //buildingUnit exists, so we've started actual construction
		{
			startedCount[b->type]++;
			b->started=true;
		}
		if (b->type.isAddon()) //if the type is an addon
		{
			if (b->builderUnit!=NULL)
				b->buildingUnit=b->builderUnit->getAddon(); //set buildingUnit to the addon (whether the addon exists or not)

			BWAPI::Unit* u = b->builderUnit;
			BWAPI::Unit* s = b->buildingUnit;
			if (s!=NULL && s->isCompleted()) //if the building exists and is completed
			{
				startedCount[b->type]--;
				plannedCount[b->type]--;
				if (u != NULL)
				{
					this->builders.erase(u);
					arbitrator->removeBid(this,u);
				}
				this->placer->freeTiles(b->tilePosition, 4,3);
				this->placer->freeTiles(b->tilePosition+BWAPI::TilePosition(4,1), 2,2);

				//If the building is complete, we can forget about it.
				this->incompleteBuildings.erase(i);
			}
			else //building either does not exist, or is incomplete
			{
				if (BWAPI::Broodwar->canMake(NULL,b->type))
				{
					if (u == NULL) //if we don't have a builder, ask for one and wait for it to be offered
						buildingsNeedingBuilders[b->type.whatBuilds().first].insert(b);
					else
					{
						if (u->getAddon() == NULL) //if the addon does not exist yet
						{
							if (u->isLifted()) //if the builder is lifted, find somehwere to land
							{
								if (BWAPI::Broodwar->getFrameCount()>b->lastOrderFrame+BWAPI::Broodwar->getLatency()*2)
								{
									if (!this->placer->canBuildHereWithSpace(b->tilePosition,b->type))
									{
										this->placer->freeTiles(b->tilePosition, 4,3);
										this->placer->freeTiles(b->tilePosition+BWAPI::TilePosition(4,1), 2,2);
										b->tilePosition=placer->getBuildLocationNear(b->tilePosition,b->type.whatBuilds().first);
										if (b->tilePosition == BWAPI::TilePositions::None)
											b->tilePosition=placer->getBuildLocationNear(b->tilePosition,b->type.whatBuilds().first,0);
										this->placer->reserveTiles(b->tilePosition, 4,3);
										this->placer->reserveTiles(b->tilePosition+BWAPI::TilePosition(4,1), 2,2);
									}
																
									u->land(b->tilePosition);
									b->lastOrderFrame = BWAPI::Broodwar->getFrameCount();
								}
							}
							else
							{
								if (u->isTraining() ) //if the builder is training, cancel it - we have control over the builder
								{
									if (u->getType()!=UnitTypes::Terran_Command_Center)
										u->cancelTrain();
								}
								else
								{
									if (u->getTilePosition()!=b->tilePosition) //we may need to lift to find a build site with enough space for an add-on
									{ //b->tilePosition is set in onOffer()
										if (u->getType()!=UnitTypes::Terran_Command_Center){
											u->lift();
											b->lastOrderFrame = BWAPI::Broodwar->getFrameCount();
										}									
									}
									else //otherwise, we may be in a good spot to make the add-on
									{
										//check to see if we can build here
										bool buildable=true;
										for(int x=b->tilePosition.x()+4;x<b->tilePosition.x()+6;x++)
											for(int y=b->tilePosition.y()+1;y<b->tilePosition.y()+3;y++)
												if (!this->placer->buildable(x,y) || BWAPI::Broodwar->hasCreep(x,y))
													buildable=false;
										if (buildable) //if so, start making the add-on
										{
											u->buildAddon(b->type);
										}
										else
										{
											if (u->getType()!=UnitTypes::Terran_Command_Center){
												//otherwise, lift up (once we are lifted we will look for a valid spot to land and build)
												u->lift();
												b->lastOrderFrame = BWAPI::Broodwar->getFrameCount();
											}						
										}
									}
								}
							}
						}
					}
				}
			}
		}
		else //if the type is a normal building (not an addon)
		{
			if (b->tilePosition==BWAPI::TilePositions::None) //if we don't have a build location yet
			{
				if ((BWAPI::Broodwar->getFrameCount()+index)%25==0)
				{
					//get a build location near the goal position
					b->tilePosition = this->placer->getBuildLocationNear(b->goalPosition, b->type);
					if (b->tilePosition!=BWAPI::TilePositions::None)
					{
						b->position = BWAPI::Position(b->tilePosition.x()*32 + b->type.tileWidth()*16, b->tilePosition.y()*32 + b->type.tileHeight()*16);
						this->placer->reserveTiles(b->tilePosition, b->type.tileWidth(), b->type.tileHeight());
					}
				}
				//if we didn't find a build location yet, or cannot make the building, skip this building for now
				if (b->tilePosition==BWAPI::TilePositions::None)
					continue;
			}

			//if the builder dies, set it to null
			if (b->builderUnit!=NULL && !b->builderUnit->exists())
				b->builderUnit=NULL;

			//if the building dies, or isn't the right type, set it to null
			if (b->buildingUnit!=NULL && (!b->buildingUnit->exists() || b->buildingUnit->getType()!=b->type))
				b->buildingUnit=NULL;

			if (b->buildingUnit == NULL) //if we don't have a building yet, look for it
			{
				//look at the units on the tile to see if it exists yet
				std::set<BWAPI::Unit*> unitsOnTile = BWAPI::Broodwar->getUnitsOnTile(b->tilePosition.x(), b->tilePosition.y());
				for(std::set<BWAPI::Unit*>::iterator t = unitsOnTile.begin(); t != unitsOnTile.end(); t++)
					if ((*t)->getType() == b->type && !(*t)->isLifted())
					{
						//we found the building
						b->buildingUnit = *t;
						break;
					}
					//maybe the builder _is_ the building! (Zerg)
					if (b->buildingUnit == NULL && b->builderUnit!=NULL && b->builderUnit->getType().isBuilding())
					{
						//we found the building
						b->buildingUnit = b->builderUnit;
					}
			}
			BWAPI::Unit* u = b->builderUnit;
			BWAPI::Unit* s = b->buildingUnit;
			if (s != NULL && s->isCompleted()) //if the building is completed, we're done
			{
				startedCount[b->type]--;
				plannedCount[b->type]--;
				if (u != NULL)
				{
					this->builders.erase(u);
					arbitrator->removeBid(this,u);
				}
				this->placer->freeTiles(b->tilePosition, b->type.tileWidth(), b->type.tileHeight());

				//If the building is complete, we can forget about it.
				this->incompleteBuildings.erase(i);
			}
			else //otherwise, it doesn't exist or is incomplete
			{
				if (s == NULL) //if the building doesn't even exist
				{
					if (u == NULL) //ask for a builder if we don't have one yet
						buildingsNeedingBuilders[b->type.whatBuilds().first].insert(b);
					else //if we have a worker
					{
						
						if (!u->isConstructing()) //if the worker isn't constructing
						{
							//add in group
							this->worker->constructingSCV.insert(u);

							double distance = u->getPosition().getDistance(b->position);
							if (distance > 64 && u->getLastCommandFrame() + 4 < BWAPI::Broodwar->getFrameCount() && u->getOrder() != BWAPI::Orders::Move ) //if its too far away, tell it to go to the build site
								u->rightClick(b->position);
							else //if its close enough, tell it to build
								if (BWAPI::Broodwar->canBuildHere(u, b->tilePosition, b->type)) //if we can build here, tell the worker to build
								{
									if (BWAPI::Broodwar->canMake(u, b->type))
									{
										u->build(b->tilePosition, b->type);
									}
								}
								else //if we cannot build here, we need to find another build site (reset the tilePosition)
								{
									//we want to build command center is built on right position, not anywhere
									if (b->type == UnitTypes::Terran_Command_Center)
									{
										//get unit set on tile
										//std::set<Unit*> unitsOnSpot = Broodwar->getUnitsInRadius(b->position,32*3);
										std::set<Unit*> unitsOnSpot = Broodwar->getUnitsInRectangle(Position(b->tilePosition),Position(b->tilePosition + TilePosition(b->type.tileWidth(),b->type.tileHeight())));
										for each (Unit* un  in unitsOnSpot)
										{
											//if this unit is ours
											if (un->getPlayer() == Broodwar->self())
											{
												//if this unit is not a building, then move away
												if (!un->getType().isBuilding())
												{
													//if the unit is vulture mine, then destroy it
													if(un->getType() == UnitTypes::Terran_Vulture_Spider_Mine)
													{														
														UnitGroup tempAttackGroup = SelectAll()(isCompleted)(canAttack).not(isWorker,isBuilding,Vulture_Spider_Mine);
														//if (!u->isConstructing())
														u->attack(un);
														if (!tempAttackGroup.empty())
														{
															Unit* mineCleanner = (*tempAttackGroup.begin());
															mineCleanner->attack(un);
														}
													}
													//if the unit is the builder or mining worker
													else if (un == b->builderUnit || (un->getType().isWorker() && (un->isGatheringGas() || un->isGatheringMinerals())))
														continue;
													//if is sieged tank
													else if (un->getType() == UnitTypes::Terran_Siege_Tank_Siege_Mode)
														un->unsiege();													
													else
													{
														un->move(BWTA::getStartLocation(Broodwar->self())->getPosition());									
													}
												}
												//if this unit is a building
												else
												{
													//unfinished unit
													if (!un->isCompleted() && un->getType()!=UnitTypes::Terran_Command_Center && !un->getType().isAddon()
														&& un->getType()!=UnitTypes::Terran_Refinery)
														un->cancelConstruction();
													//for a completed building
													else if (un->isCompleted()){
														//if this unit can lift, then return
														if (un->getType()!=UnitTypes::Terran_Command_Center && un->lift())															
															continue;
	
														else if (un->getType()==UnitTypes::Terran_Command_Center || un->getType().isAddon()
															|| un->getType()==UnitTypes::Terran_Refinery)
															continue;
														//if the building can not lift,then destroy it
														else{
															UnitGroup tempAttackGroup = SelectAll()(isCompleted)(canAttack).not(isWorker,isBuilding,Vulture_Spider_Mine);
															if (Broodwar->getFrameCount()%24==0){
                                WorkerManager::create()->addToNotRepairList(un);
																if(!u->isConstructing())
																	u->attack(un);
																tempAttackGroup.attackUnit(un);
																return;
															}
														}
													}
												}
											}
											else
											{
												//if it is a neutral building or unit
												if (un->getPlayer()->isNeutral())
												{
													if (un->isInvincible())
													{
														this->placer->freeTiles(b->tilePosition, b->type.tileWidth(), b->type.tileHeight());
														b->tilePosition = BWAPI::TilePositions::None;
														b->position = BWAPI::Positions::None;
														break;
													}
													else
													{
														UnitGroup tempAttackGroup = SelectAll()(isCompleted)(canAttack).not(isWorker,isBuilding);
														if (Broodwar->getFrameCount()%24==0)
														{
															if(!u->isConstructing())
																u->attack(un);
															tempAttackGroup.attackUnit(un);
															return;
														}
													}
												}
												//if it is enemy unit
												else
												{
													// invisible units: dark templar...
													if (!un->isDetected())
													{
														UnitGroup scanners = SelectAll()(isCompleted)(Comsat_Station)(Energy,">=",50);
														if (!scanners.empty())
														{
															Unit* scanner = *(scanners.begin());
															scanner->useTech(TechTypes::Scanner_Sweep,un->getPosition());
														}
													}
													else
													{
														UnitGroup tempAttackGroup = SelectAll()(isCompleted)(canAttack).not(isWorker,isBuilding);
														if (Broodwar->getFrameCount()%24 == 0)
														{
															//if(!u->isConstructing())
															u->attack(un);
															tempAttackGroup.attackUnit(un);
															return;
														}
													}												
												}
												//this->placer->freeTiles(b->tilePosition, b->type.tileWidth(), b->type.tileHeight());
												//b->tilePosition = BWAPI::TilePositions::None;
												//b->position = BWAPI::Positions::None;
											}								
										}
									}
									this->placer->freeTiles(b->tilePosition, b->type.tileWidth(), b->type.tileHeight());
									b->tilePosition = BWAPI::TilePositions::None;
									b->position = BWAPI::Positions::None;
								}              
						}
					}
				}
				else //if the building does exist
				{
					if (b->type.getRace() != BWAPI::Races::Terran) //if its not terran, we can forget about the builder (i.e. probe is not needed to finish construction)
					{
						if (u != NULL)
						{
							this->builders.erase(u);
							arbitrator->removeBid(this, u);
							b->builderUnit = NULL;  
							u = b->builderUnit;
						}
					}
					else
					{
						//if the building is terran, the worker may have been killed
						if (u == NULL) //looks like the worker was killed, or revoked. In either case we need to ask for another worker to finish our building
							buildingsNeedingBuilders[b->type.whatBuilds().first].insert(b);
            else if (u->getHitPoints() < 15) // if dont have enough hp, remove
            {
              this->builders.erase(u);
              arbitrator->removeBid(this, u);
              u->haltConstruction();
              u = b->builderUnit = NULL;
            }
						else
						{
							//we have a worker, so lets rightClick it on the incomplete building so it can resume construction
							if (BWAPI::Broodwar->getFrameCount()%(4*BWAPI::Broodwar->getLatency())==0)
							{
								if (( !u->isConstructing()) || !s->isBeingConstructed() )
								{
									//right click builder on building
									u->rightClick(s);
								}
							}
						}
					}
				}
			}
		}
	}
}

std::string ConstructionManager::getName() const
{
	return "Construction Manager";
}

std::string ConstructionManager::getShortName() const
{
	return "Con";
}

void ConstructionManager::onRemoveUnit(BWAPI::Unit* unit)
{
	//remove the builder if needed - incomplete buildings will ask for a new builder during update()
	if (builders.find(unit) != builders.end())
	{
		Building* building = builders.find(unit)->second;
		building->builderUnit = NULL;
		builders.erase(unit);
	}
	else
	{
		for(std::list<Building>::iterator i=incompleteBuildings.begin();i!=incompleteBuildings.end();i++)
		{
			if (unit==i->buildingUnit)
			{
				i->buildingUnit=NULL;
			}
		}
	}
}

bool ConstructionManager::build(BWAPI::UnitType type, BWAPI::TilePosition goalPosition)
{
	//build order starts here
	if (!type.isBuilding()) return false; //we only accept buildings
	Building newBuilding;
	newBuilding.type           = type;
	newBuilding.goalPosition   = goalPosition;
	newBuilding.tilePosition   = BWAPI::TilePositions::None;
	newBuilding.builderUnit    = NULL; //builder not picked yet
	newBuilding.buildingUnit   = NULL; //building does not exist yet
	newBuilding.position       = BWAPI::Positions::None; //no position yet
	newBuilding.lastOrderFrame = 0;
	newBuilding.started        = false; //not started yet
	plannedCount[type]++; //increment planned count for this type of unit
	this->incompleteBuildings.push_back(newBuilding);
	return true;
}

int ConstructionManager::getPlannedCount(BWAPI::UnitType type) const
{
	std::map<BWAPI::UnitType,int>::const_iterator i=plannedCount.find(type);
	if (i!=plannedCount.end())
		return i->second;
	return 0;
}

void ConstructionManager::deleteBuilding(BWAPI::UnitType type, BWTA::Region* r)
{
  if (!type.isBuilding()) return;
  // Remove building from incompleteBuildings list
  for (std::list<Building>::iterator i = incompleteBuildings.begin(), j; i!=incompleteBuildings.end();){
    j = i++;
    Building* b = &(*(j));
    if (b->type == type &&
        (!r || BWTA::getRegion(b->tilePosition)==r))
    {
      //Broodwar->printf("Plan %s in position (%d,%d) is removed", b->type.c_str(),b->position.x(),b->position.y());
      Unit* u = b->builderUnit;
      startedCount[b->type]--;
      plannedCount[b->type]--;
      // if there is builder, halt construction and remove from bid
      if (u != NULL)
      {
        if (!type.isAddon())
        {
          u->haltConstruction();
        }
        this->builders.erase(u);
        arbitrator->removeBid(this,u);
      }
      // if we already build this building, cancel it
      if (b->buildingUnit)
      {
        b->buildingUnit->cancelConstruction();
      }
      this->placer->freeTiles(b->tilePosition, b->type.tileWidth(), b->type.tileHeight());

      // this building dont need anymore worker
      buildingsNeedingBuilders[b->type.whatBuilds().first].erase(b);

      // remove the building
      this->incompleteBuildings.erase(j);
    }
  }
}

int ConstructionManager::getStartedCount(BWAPI::UnitType type) const
{
	std::map<BWAPI::UnitType,int>::const_iterator i=startedCount.find(type);
	if (i!=startedCount.end())
		return i->second;
	return 0;
}